#黄金正弦算法
import numpy as np
import math
import copy
import time

class GSA():
    def __init__(self, n_dim=None, pop_size=20, max_iter=50, lb=-512, ub=512, obj_func=None):
        self.dim = n_dim
        self.pop = pop_size
        self.ub = ub
        self.lb = lb
        self.Max_iter = max_iter
        self.fun = obj_func
        self.X = np.zeros([pop_size, n_dim])
        for i in range(pop_size):
            for j in range(n_dim):
                self.X[i, j] = (ub[j] - lb[j]) * np.random.random() + lb[j]

    def BorderCheck(self):
        #边界检查
        pop = self.pop
        dim = self.dim
        lb = self.lb
        ub = self.ub
        for i in range(pop):
            for j in range(dim):
                if self.X[i,j] > ub[j]:
                    self.X[i, j] = ub[j]
                elif self.X[i,j] < lb[j]:
                    self.X[i, j] = lb[j]

    '''计算适应度函数'''

    def CaculateFitness(self):
        pop = self.X.shape[0]
        fitness = np.zeros([pop, 1])
        for i in range(pop):
            fitness[i] = self.fun(self.X[i, :])
        return fitness

    '''对适应度排序'''

    def SortFitness(self, Fit):
        fitness = np.sort(Fit, axis=0)
        index = np.argsort(Fit, axis=0)
        return fitness, index

    '''根据适应度对位置进行排序'''

    def SortPosition(self, index):
        Xnew = np.zeros(self.X.shape)
        for i in range(self.X.shape[0]):
            Xnew[i, :] = self.X[index[i], :]
        return Xnew

    def Tent_GSA(self):
        pop = self.pop
        dim = self.dim
        ub = self.ub
        lb = self.lb
        fun = self.fun
        maxIter = self.Max_iter
        a = -math.pi
        b = math.pi
        gold = (np.sqrt(5) - 1)/2  #黄金分割率
        x1 = a+(1-gold)*(b-a)      #黄金分割系数
        x2 = a+gold*(b-a)          #黄金分割系数
        fitness = self.CaculateFitness()
        fitness, sortIndex = self.SortFitness(fitness)
        self.X = self.SortPosition(sortIndex)
        GbestScore = copy.copy(fitness[0])
        GbestPositon = np.zeros([1, dim])
        GbestPositon[0, :] = copy.copy(self.X[0, :])
        Curve = np.zeros([maxIter, 1])
        avg_fit = []  # 每次迭代的所有个体适应度平均值
        avg_ts = []  # 平均方差
        time_start = time.time()  # 记录迭代寻优开始时间
        for t in range(maxIter):
            for i in range(pop):
                r = np.random.random()
                r1 = 2*math.pi*r
                r2 = r*math.pi
                for j in range(dim):
                    self.X[i,j] = self.X[i,j]*np.abs(np.sin(r1)) - r2*np.sin(r1)*np.abs(x1*GbestPositon[0,j] - x2*self.X[i,j])

            self.BorderCheck()
            fitness = self.CaculateFitness()
            avg_fit.append(np.mean(fitness))
            avg_ts.append(np.var(fitness))
            #更新黄金分割系数
            for i in range(pop):
                if fitness[i] < GbestScore:    # 当前解优于最优解
                    GbestScore = fitness[i]
                    GbestPositon[0,:] = copy.copy(self.X[i,:])
                    b = x2
                    x2 = x1
                    x1 = a+(1-gold)*(b-a)
                else:
                    a = x1
                    x1 = x2
                    x2 = a+gold*(b-a)

                if x1 == x2:     #若黄金分割系数相同，则随机重置黄金分割系数
                    a = -math.pi*np.random.random()
                    b = math.pi * np.random.random()
                    x1 = a+(1-gold)*(b-a)
                    x2 = a+gold*(b-a)

            fitness,sortIndex = self.SortFitness(fitness)
            self.X = self.SortPosition(sortIndex)
            if fitness[0] <= GbestScore:
                GbestScore = copy.copy(fitness[0])
                GbestPositon[0, :] = copy.copy(self.X[0, :])
            Curve[t] = GbestScore
        time_end = time.time()  # 记录迭代结束时间
        print(f'GSA共花费 {time_end - time_start} 秒')

        return Curve, avg_fit, avg_ts
